package promote.brozen.algorithm.tree.balance.bt;

import promote.brozen.algorithm.tree.Tree;
import promote.brozen.utils.Assert;

import java.util.function.Consumer;

/**
 * @author Brozen
 * @date 2020/1/13 4:21 PM
 */
public class BTree<K extends Comparable<K>> implements Tree<K> {

    /**
     * B-树的阶
     */
    private final int degree;

    private BTreeNode<K> root;

    public BTree() {
        this(3);
    }

    public BTree(int degree) {
        Assert.isTrue(degree >= 2, "B-Tree's degree must greater than or equals to 2!");
        this.degree = degree;
    }

    @Override
    public void add(K key) {
        // 如果根节点为空，则key设置为根节点
        if (this.root == null) {
            this.root = new BTreeNode<>(key);
            return;
        }

        // 1. 执行插入操作
        BTreeDataChainNode<K> chainNode = doAdd(this.root, key);

        // 2. 如果插入后，新增key所在节点的keys数量超过degree-1，需要执行平衡
        if (chainNode.getNode().getKeySize() >= this.degree) {
            this.doAddBalance(chainNode.getNode());
        }
    }

    private BTreeDataChainNode<K> doAdd(BTreeNode<K> treeNode, K key) {
        BTreeDataChainNode<K> chainNode = treeNode.getHead().getNext();
        int result;
        // 1. 循环树节点chain，找到要插入到chain中的位置
        while (chainNode != null) {
            result = chainNode.getKey().compareTo(key);
            if (result < 0) {
                chainNode = chainNode.getNext().getNext();
            } else if (result > 0) {
                // 1.2 找到了要插入chain节点的后面的chain节点
                break;
            } else {
                // 1.1 存在一个key相同的节点，需要更新
                chainNode.setKey(key);
                return chainNode;
            }
        }

        // 2. 尝试插入节点
        if (treeNode.isLeaf()) {
            // 2.2 当前树节点是叶子节点，需要插入到当前位置
            BTreeDataChainNode<K> chainDataNode = new BTreeDataChainNode<>(treeNode, key);
            BTreeChildChainNode<K> chainChildNodeNext = new BTreeChildChainNode<>(treeNode);
            chainDataNode.setNext(chainChildNodeNext);
            chainChildNodeNext.setPrev(chainDataNode);

            if (chainNode != null) {
                // 2.1.1 插入chain节点的后一个chain节点存在，则说明是插入到链表中间
                chainDataNode.setPrev(chainNode.getPrev());
                chainNode.getPrev().setNext(chainDataNode);
                chainNode.setPrev(chainChildNodeNext);
                chainChildNodeNext.setNext(chainNode);
            } else {
                // 2.1.2 后一个chain节点不存在，说明插入到链表的尾部
                BTreeChildChainNode<K> tail = treeNode.getTail();
                tail.setNext(chainDataNode);
                chainDataNode.setPrev(tail);
                treeNode.setTail(chainChildNodeNext);
            }
            treeNode.setKeySize(treeNode.getKeySize() + 1);
            return chainDataNode;
        } else {
            // 2.1 当前树节点不是叶子节点，需要向子节点中插入key
            if (chainNode == null) {
                return doAdd(treeNode.getTail().getChild(), key);
            } else {
                return doAdd(chainNode.getPrev().getChild(), key);
            }
        }
    }

    /**
     * 进入平衡操作的平衡节点，相同点在于，平衡节点包含key的个数大于或等于degree；
     * pi为平衡节点在其父节点中指针在children中的槽位
     */
    private void doAddBalance(BTreeNode<K> node) {
        if (node == null || node.getKeySize() < this.degree) {
            return;
        }

        // 1. 如果node等于root节点，需要将平衡节点裂开为两个
        // 1.1 首先找到ki，i是keys中间的槽位
        // 1.2 将keys中，小于i的槽位保持不动，大于i的槽位新建一个节点存储
        // 1.3 新建一个节点，将ki放入，并设置c0指向平衡节点、c1指向新增的节点k i+1 ~ kn

        // 1. 找到中间位置的key
        int oldKeySize = node.getKeySize();
        int i = oldKeySize % 2 == 0 ? oldKeySize / 2 - 1 : oldKeySize / 2;
        BTreeDataChainNode<K> ki = node.getDataChainNode(i);

        // 2 将node从中间位置裂开
        // 2.1 小于中间位置的节点保持不变
        BTreeChildChainNode<K> oldTail = node.getTail();
        node.setTail(ki.getPrev());
        ki.getPrev().setNext(null);
        node.setKeySize(i);
        // 2.2 大于中间位置的节点用新的节点存储
        BTreeNode<K> newTreeNode = new BTreeNode<>();
        BTreeChildChainNode<K> newHead = ki.getNext();
        newTreeNode.setHead(newHead);
        newHead.setPrev(null);
        newTreeNode.setTail(oldTail);
        newTreeNode.setKeySize(oldKeySize - i - 1);
        newTreeNode.setLeaf(node.isLeaf());
        // 2.2.1 设置chain节点所属的树节点
        BTreeChildChainNode<K> n = newHead;
        while (n != null) {
            n.setNode(newTreeNode);
            if (n.getNext() != null) {
                n.getNext().setNode(newTreeNode);
                n = n.getNext().getNext();
            } else {
                n = null;
            }
        }

        // 3 找到父节点中指向平衡节点的指针
        BTreeChildChainNode<K> parentChainNode = node.getParent();
        if (parentChainNode == null) {
            // 3.1 如果这个指针不存在，说明是根节点裂开，需要新建节点，此时树会长高，也是BTree长高的唯一途径：根节点裂开
            // 3.1.1 新建一个根节点
            BTreeNode<K> newRoot = new BTreeNode<>();
            newRoot.setLeaf(false);
            // 3.1.2 head指向平衡节点(也就是根节点)，tail指向分裂后的新子节点
            BTreeChildChainNode<K> newRootHead = newRoot.getHead();
            BTreeChildChainNode<K> newRootTail = newRoot.getTail();
            newRootHead.setChild(node);
            node.setParent(newRootHead);
            newRootTail.setChild(newTreeNode);
            newTreeNode.setParent(newRootTail);
            // 3.1.3 把ki放到新的根节点中
            ki.setNode(newRoot);
            newRootHead.setNext(ki);
            ki.setPrev(newRootHead);
            ki.setNext(newRootTail);
            newRootTail.setPrev(ki);
            newRoot.setKeySize(1);
            this.root = newRoot;
        } else {
            // 3.2 如果这个指针存在，则是非根节点裂开
            // 3.2.1 找到父节点
            BTreeNode<K> parentTreeNode = parentChainNode.getNode();
            ki.setNode(parentTreeNode);
            // 3.2.2 新建一个子节点指针，放在ki后面，指向分裂时新产生的节点
            BTreeChildChainNode<K> ci = new BTreeChildChainNode<>(parentTreeNode);
            ki.setNext(ci);
            ci.setPrev(ki);
            ci.setChild(newTreeNode);
            newTreeNode.setParent(ci);
            // 3.2.3 将ki和ci放到父节点中去，需要考虑在父节点中指向平衡节点的指针是不是父节点的tail
            ki.setPrev(parentChainNode);
            ci.setNext(parentChainNode.getNext());
            if (parentChainNode.getNext() != null) {
                // 3.2.3.1 在父节点中的指针的next不等于null，说明不是tail节点，需要设置后面节点的指向
                parentChainNode.getNext().setPrev(ci);
                parentChainNode.setNext(ki);
            } else {
                // 3.2.3.2 next为null，说明在父节点中的指针是tail，此时需要重设tail节点
                parentChainNode.setNext(ki);
                parentTreeNode.setTail(ci);
            }

            // 3.2.4 如果父节点中节点数量超过了degree - 1，需要将父节点带入平衡操作递归处理
            parentTreeNode.setKeySize(parentTreeNode.getKeySize() + 1);
            if (parentTreeNode.getKeySize() >= this.degree) {
                doAddBalance(parentTreeNode);
            }
        }
    }

    @Override
    public void remove(K key) {
        // 1. 找到要删除的节点
        BTreeDataChainNode<K> toDeleteDataNode = searchChainNode(this.root, key);
        if (toDeleteDataNode == null) {
            return;
        }

        // 3. 如果删除的数据节点不在叶子节点上
        if (!toDeleteDataNode.getNode().isLeaf()) {
            // 3.1 找到删除节点的前驱节点(比删除节点小的最大节点，删除节点prev指向子树的最大值)
            BTreeDataChainNode<K> predecessor = findMax(toDeleteDataNode.getPrev().getChild());
            // 3.2 用前驱节点中的值代替删除节点中的值，然后将前驱节点当做删除节点处理
            // 前驱节点必定是叶子节点
            if (predecessor != null) {
                toDeleteDataNode.setKey(predecessor.getKey());
                toDeleteDataNode = predecessor;
            }
        }

        // 2. 如果删除的数据节点是在叶子节点上
        // 2.1 执行删除
        BTreeNode<K> treeNode = toDeleteDataNode.getNode();
        if (toDeleteDataNode.getPrev() == treeNode.getHead() && toDeleteDataNode.getNext() == treeNode.getTail()) {
            // 2.1.1 如果删除节点的前指针是head，后指针是tail，说明是这个树节点的唯一key节点，需要把指针置空
            toDeleteDataNode.getPrev().setNext(null);
            toDeleteDataNode.getNext().setPrev(null);
            toDeleteDataNode.setNext(null);
            toDeleteDataNode.setPrev(null);
        } else if (toDeleteDataNode.getPrev() == treeNode.getHead()) {
            // 2.1.2 删除节点的前指针是head，后指针不是tail，说明删除节点是链表中第一个，后面还有其他节点，需要重置head指针
            treeNode.setHead(toDeleteDataNode.getNext());
            toDeleteDataNode.getNext().setPrev(null);
            toDeleteDataNode.setNext(null);
        } else if (toDeleteDataNode.getNext() == treeNode.getTail()) {
            // 2.1.3 删除节点的前指针不是head，后指针是tail，说明删除节点是链表中最后一个，前面有其他节点，需要重置tail指针
            treeNode.setTail(toDeleteDataNode.getPrev());
            toDeleteDataNode.getPrev().setNext(null);
            toDeleteDataNode.setPrev(null);
        } else {
            // 2.1.4 删除节点不是第一个也不是最后一个，在中间，则设置前后指针即可
            toDeleteDataNode.getPrev().setNext(toDeleteDataNode.getNext().getNext());
            toDeleteDataNode.getNext().getNext().setPrev(toDeleteDataNode.getPrev());
        }
        treeNode.setKeySize(treeNode.getKeySize() - 1);

        // 2.2 如果删除后叶子节点为空，需要进行平衡操作
        // 其实只有2.1.1中的情况会导致树节点为空
        if (toDeleteDataNode.getNode().getKeySize() == 0) {
            doRemoveBalance(treeNode);

        }
    }

    /**
     * 删除时的平衡操作
     * 递归进入条件：
     *  需要进行平衡操作的树节点都是key为空的节点
     *
     * 可以推导出：平衡节点的头尾都指向同一个节点
     *
     * 递归临界条件：
     *  1. 平衡节点为根节点
     *  2. 平衡完成后，平衡节点父节点不为空
     *
     */
    private void doRemoveBalance(BTreeNode<K> treeNode) {
        // 0. 如果平衡节点为根节点，终止平衡，重设根节点为平衡节点的子节点

        // 1. 如果平衡节点的兄弟节点中的key足够多，从兄弟节点借一个过来，完成平衡
        BTreeChildChainNode<K> parentPointer = treeNode.getParent();
        BTreeNode<K> parentTreeNode = parentPointer.getNode();
        BTreeChildChainNode<K> leftSiblingPointer = parentPointer != parentTreeNode.getHead() ? parentPointer.getPrev().getPrev() : null;
        BTreeNode<K> leftSibling = leftSiblingPointer == null ? null : leftSiblingPointer.getChild();
        BTreeChildChainNode<K> rightSiblingPointer = parentPointer != parentTreeNode.getTail() ? parentPointer.getNext().getNext() : null;
        BTreeNode<K> rightSibling = rightSiblingPointer == null ? null : rightSiblingPointer.getChild();
        if (leftSibling != null && leftSibling.getKeySize() > 1) {
            // 1.1 如果左兄弟存在，且左兄弟key节点多于1个，则从左兄弟借一个
            // 1.1.1 左兄弟的最后一个key和tail拿出来，需要重设左兄弟tail
            BTreeDataChainNode<K> borrowed = leftSibling.getTail().getPrev();
            leftSibling.setTail(borrowed.getPrev());
            leftSibling.getTail().setNext(null);
            leftSibling.setKeySize(leftSibling.getKeySize() - 1);
            // 1.1.2 借出来的key节点放到平衡节点中来：并用借来节点的next作为平衡节点新的head
            //       借来的节点放在平衡节点的head和tail中间
            BTreeChildChainNode<K> newHead = borrowed.getNext();
            BTreeChildChainNode<K> tail = treeNode.getTail();
            newHead.setNext(borrowed);
            newHead.setPrev(null);
            borrowed.setPrev(newHead);
            borrowed.setNext(tail);
            tail.setPrev(borrowed);
            treeNode.setHead(newHead);
            borrowed.setNode(treeNode);
            newHead.setNode(treeNode);
            // 1.1.3 交换父节点指针的prev与借来节点的key值
            BTreeDataChainNode<K> parentPointerPrev = parentPointer.getPrev();
            K borrowedKey = parentPointerPrev.getKey();
            parentPointerPrev.setKey(borrowed.getKey());
            borrowed.setKey(borrowedKey);
        } else if (rightSibling != null && rightSibling.getKeySize() > 1) {
            // 1.2 如果右兄弟存在，且右兄弟key节点多于1个，则从右兄弟借一个
            // 1.2.1 右兄弟的第一个key和head取出来，重设右兄弟head
            BTreeDataChainNode<K> borrowed = rightSibling.getHead().getNext();
            rightSibling.setHead(borrowed.getNext());
            rightSibling.getHead().setPrev(null);
            rightSibling.setKeySize(rightSibling.getKeySize() - 1);
            // 1.2.2 用借来节点的prev代替平衡节点的tail，并将借来的节点与其prev交换位置，借来节点的prev指向平衡节点的head
            BTreeChildChainNode<K> head = treeNode.getHead();
            BTreeChildChainNode<K> newTail = borrowed.getPrev();
            newTail.setPrev(borrowed);
            newTail.setNext(null);
            borrowed.setNext(newTail);
            borrowed.setPrev(head);
            head.setNext(borrowed);
            treeNode.setTail(newTail);
            borrowed.setNode(treeNode);
            newTail.setNode(treeNode);
            // 1.2.3 交换父节点指针的next与借来节点的key值
            BTreeDataChainNode<K> parentPointerNext = parentPointer.getNext();
            K borrowedKey = parentPointerNext.getKey();
            parentPointerNext.setKey(borrowed.getKey());
            borrowed.setKey(borrowedKey);
        } else {
            // 2. 如果兄弟节点key不够，则将平衡节点与兄弟节点合并，并从父节点降一个节点下来，作为与兄弟节点合并的中间节点；
            if (leftSibling != null) {
                // 2.1 若左兄弟存在，优先与左兄弟合并，让父节点指针的prev降下来；
                // 2.1.1 把下降的key节点从父节点中移除掉
                BTreeDataChainNode<K> sunk = parentPointer.getPrev();
                BTreeChildChainNode<K> sunkPrev = sunk.getPrev();
                BTreeDataChainNode<K> sunkNext = sunk.getNext().getNext();
                sunkPrev.setNext(sunkNext);
                if (sunkNext != null) {
                    sunkNext.setPrev(sunkPrev);
                } else {// sunkNext为空，说明sunk.next是尾节点，sunk是最后一个key，此时需要重新设置tail
                    parentTreeNode.setTail(sunk.getPrev());
                }
                parentTreeNode.setKeySize(parentTreeNode.getKeySize() - 1);
                // 2.1.2 把下降节点拼接到左兄弟的tail上
                BTreeChildChainNode<K> newTail = sunk.getNext();
                BTreeChildChainNode<K> oldTail = leftSibling.getTail();
                sunk.setPrev(oldTail);
                oldTail.setNext(sunkNext);
                leftSibling.setTail(newTail);
                sunk.setNode(leftSibling);
                newTail.setNode(leftSibling);
                leftSibling.setKeySize(leftSibling.getKeySize() + 1);
                // 2.1.3 设置新的左兄弟的tail的child指针指向平衡节点的head的child
                newTail.setChild(treeNode.getHead().getChild());
            } else if (rightSibling != null) {
                // 2.2 若右兄弟存在，与右兄弟合并，让父节点指针的next降下来；
                // 2.2.1 把下降节点从父节点中移除
                BTreeDataChainNode<K> sunk = parentPointer.getNext();
                BTreeDataChainNode<K> sunkPrev = sunk.getPrev().getPrev();
                BTreeChildChainNode<K> sunkNext = sunk.getNext();
                sunkNext.setPrev(sunkPrev);
                if (sunkPrev != null) {
                    sunkPrev.setNext(sunkNext);
                } else {
                    parentTreeNode.setHead(sunkNext);
                }
                parentTreeNode.setKeySize(parentTreeNode.getKeySize() - 1);
                // 2.2.2 把下降节点拼接到右兄弟的head上
                BTreeChildChainNode<K> newHead = sunk.getPrev();
                BTreeChildChainNode<K> oldHead = rightSibling.getHead();
                oldHead.setPrev(sunk);
                sunk.setNext(oldHead);
                rightSibling.setHead(newHead);
                sunk.setNode(rightSibling);
                newHead.setNode(rightSibling);
                rightSibling.setKeySize(rightSibling.getKeySize() + 1);
                // 2.1.3 设置新的右兄弟head的child指向平衡节点tail的child
                newHead.setChild(treeNode.getTail().getChild());
            } else {
                // 左右兄弟均为空，这是树出了问题
                throw new IllegalStateException("BTree has some problem: parent has only one child, parent=" + parentTreeNode.toString());
            }
        }

        if (parentTreeNode.getKeySize() <= 0) {
            // 3. 如果平衡节点与兄弟节点合并后，父节点为空，则将父节点带入平衡操作递归处理；
            doRemoveBalance(parentTreeNode);
        }
    }

    private BTreeDataChainNode<K> findMax(BTreeNode<K> node) {
        if (node == null) {
            return null;
        }

        BTreeChildChainNode<K> tail = node.getTail();
        if (tail.getChild() != null) {
            return findMax(tail.getChild());
        } else {
            return tail.getPrev();
        }
    }

    @Override
    public K search(K key) {
        BTreeDataChainNode<K> node = searchChainNode(this.root, key);
        return node == null ? null : node.getKey();
    }

    private BTreeDataChainNode<K> searchChainNode(BTreeNode<K> node, K key) {
        if (node == null) {
            return null;
        }

        BTreeDataChainNode<K> item = node.getHead().getNext();
        int result;
        while (item != null) {
            result = item.getKey().compareTo(key);
            if (result < 0) {
                item = item.getNext().getNext();
            } else if (result > 0) {
                return searchChainNode(item.getPrev().getChild(), key);
            } else {
                return item;
            }
        }

        if (node.isLeaf()) {
            return null;
        } else {
            return searchChainNode(node.getTail().getChild(), key);
        }
    }

    public void walk(Consumer<K> consumer) {
        doWalk(this.root, consumer);
    }

    private void doWalk(BTreeNode<K> treeNode, Consumer<K> consumer) {
        BTreeChildChainNode<K> node = treeNode.getHead();
        while (node != null) {
            // 存在子节点，先遍历子节点
            if (node.getChild() != null) {
                doWalk(node.getChild(), consumer);
            }

            if (node.getNext() != null) {
                consumer.accept(node.getNext().getKey());
                node = node.getNext().getNext();
                continue;
            }

            break;
        }
    }
}
